import cv2
import cv2.ximgproc as ximgproc
import numpy as np
import math
import os
import time
import tools as f


class RebarDetector:
    """
    钢筋检测主类，实现钢筋检测全流程
    采用模板方法模式定义检测流程
    """
    def __init__(self, model_path="model.yml"):
        self.isShow = False  # 是否显示中间结果
        self.edge_detector = ximgproc.createStructuredEdgeDetection(model_path)
        self.lsd_detector = cv2.createLineSegmentDetector(cv2.LSD_REFINE_STD)
        # 初始化各个处理模块
        self.line_classifier = f.LineClassifier() # 线段分类器
        self.line_fitter = f.LineFitter()   # 线段拟合器
        self.intersection_finder = f.IntersectionFinder() # 交点查找器
        self.mask_creator = f.MaskCreator() # 掩膜创建器
        if self.isShow:
            self.result_visualizer = f.ResultVisualizer()   # 结果可视化器
        
    def process_image(self, ir_img, depth_img):
        """
        处理图像
        """
        if ir_img is None or depth_img is None:
            return None
        s1 = time.time()
        # 边缘检测
        edges = self._detect_edges(depth_img)
        s2 = time.time()
        # print(f"边缘检测耗时: {s2 - s1:.2f} 秒")
        # 直线检测
        lines = self._detect_lines(edges)
        s3 = time.time()
        # print(f"直线检测耗时: {s3 - s2:.2f} 秒")
        if len(lines) < 100:
            # print("检测到的线段数量过少，可能图像有问题")
            return None
        # 分类线段
        horizontal, vertical, avg_h_angle, avg_v_angle = self.line_classifier.classify(lines)
        s4 = time.time()
        # print(f"线段分类耗时: {s4 - s3:.2f} 秒")
        # 拟合直线
        h_fitted = self.line_fitter.fit(horizontal, avg_h_angle,depth_img)
        v_fitted = self.line_fitter.fit(vertical, avg_v_angle,depth_img)
        # v_fitted = []
        s5 = time.time()
        # print(f"线段拟合耗时: {s5 - s4:.2f} 秒")
        # 计算交点
        intersections = self.intersection_finder.find(h_fitted, v_fitted, edges.shape[1], edges.shape[0],5)
        if len(intersections) < 5:
            # print("未检测到足够的交点，可能图像有问题")
            return None
        
        s6 = time.time()
        # print(f"交点计算耗时: {s6 - s5:.2f} 秒")
        # 创建掩膜
        masked_img = self.mask_creator.create(ir_img, intersections)
        s7 = time.time()
        # print(f"掩膜创建耗时: {s7 - s6:.2f} 秒")
        # 绘制结果
        if self.isShow:
            fitted_img = self.result_visualizer.draw_lines(depth_img, h_fitted, v_fitted)
            fitted_img = self.result_visualizer.draw_intersections(fitted_img, intersections)
            result = {
                'original': ir_img,
                'edges': edges,
                'lines': lines,
                'fitted': fitted_img,
                'masked': masked_img,
                'intersections': intersections
            }
            self.result_visualizer._display_results(result)
        return masked_img
        
    def _detect_edges(self, depth_img):
        """
        边缘检测
        """
        # 下采样加速处理（保持宽高比）
        scale = 0.5  # 根据实际效果调整
        small_img = cv2.resize(depth_img, (0,0), fx=scale, fy=scale)
        src = cv2.cvtColor(small_img, cv2.COLOR_BGR2RGB)
        edges = self.edge_detector.detectEdges(np.float32(src) / 255.0)
        edges = cv2.normalize(edges, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
        # 二值化
        _, binary = cv2.threshold(edges, 10, 255, cv2.THRESH_BINARY)
        # 闭操作补缺口
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7,7))
        kernel_img = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
        # 骨架提取
        skeleton = cv2.ximgproc.thinning(kernel_img, thinningType=cv2.ximgproc.THINNING_ZHANGSUEN)
        skeleton = cv2.resize(skeleton, (depth_img.shape[1], depth_img.shape[0]))
        return skeleton
        
    def _detect_lines(self, edges):
        """
        直线检测
        """
        return self.lsd_detector.detect(edges)[0]

